/*
 *  webpage_hunter.cpp
 *  crawler
 *
 *  Created by Roy Zuo on May 13, 2010.
 *  Copyright 2010 O2 Micro, Inc. All rights reserved.
 *
 */

#include "webpage_hunter.hpp"
#include <boost/bind.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>

WebpageHunter::WebpageHunter ( boost::asio::io_service& io_service, ostream* page_header, ostream* page_content, const string& server, const string& port, const string& host, const string& path, HttpHeaderPairs& headers, signed int timeout_seconds/* =60 */ )
: _resolver ( io_service ), _socket ( io_service ), _deadtimer ( io_service ), _header_output ( page_header ), _content_output ( page_content ), _deadseconds ( timeout_seconds ), _response_stream ( 0 )
{
	ostream request ( &_request );

	generate_request ( request, host, path, headers );

	boost::asio::ip::tcp::resolver::query query ( server, port );

	_resolver.async_resolve ( query, 
		boost::bind ( &WebpageHunter::handle_resolve,
		this, 
		boost::asio::placeholders::error,
		boost::asio::placeholders::iterator ) );
}

WebpageHunter::~WebpageHunter ()
{
	if ( _response_stream )
	{
		delete _response_stream;
		_response_stream = 0;
	}

	_socket.close();
}

void WebpageHunter::generate_request (ostream& request, const string& host, const string& path, const HttpHeaderPairs& headers)
{
	request << "GET " << path << " HTTP/1.1\r\n";
	request << "Host: " << host << "\r\n";
	
	for ( HttpHeaderPairs::const_iterator iter = headers.begin(); iter != headers.end(); iter ++ )
	{
		string key = (*iter).first;
		string value = (*iter).second;
		if (value != "")
		{
			request << key << ": " << value << "\r\n";
		}
	}
	request << "\r\n";
}

void WebpageHunter::handle_resolve( const boost::system::error_code& err, boost::asio::ip::tcp::resolver::iterator endpoint_iterator )
{
	if ( err )
	{
		throw boost::system::system_error ( err ); 	
	}
	else
	{
		boost::asio::ip::tcp::endpoint endpoint = *endpoint_iterator;
		_socket.async_connect ( endpoint, 
							   boost::bind (
											&WebpageHunter::handle_connect,
											this,
											boost::asio::placeholders::error,
											++ endpoint_iterator
											) );
		
		_deadtimer.expires_from_now ( boost::posix_time::seconds ( _deadseconds ) );
		_deadtimer.async_wait ( boost::bind ( &WebpageHunter::handle_timeout, this, boost::asio::placeholders::error ) );
	}
}

void WebpageHunter::handle_connect( const boost::system::error_code& err, boost::asio::ip::tcp::resolver::iterator endpoint_iterator )
{
	if ( err )
	{
		if ( endpoint_iterator != boost::asio::ip::tcp::resolver::iterator() )
		{
			_deadtimer.cancel();	
			boost::asio::ip::tcp::endpoint  endpoint = *endpoint_iterator;
			_socket.async_connect ( endpoint, 
								   boost::bind (
												&WebpageHunter::handle_connect,
												this,
												boost::asio::placeholders::error,
												++ endpoint_iterator
												) );
			
			_deadtimer.expires_from_now ( boost::posix_time::seconds ( _deadseconds ) );
			_deadtimer.async_wait ( boost::bind ( &WebpageHunter::handle_timeout, this, boost::asio::placeholders::error ) );
		}
		else
		{
			_deadtimer.cancel();
			throw boost::system::system_error ( err ); 	
		}
	}
	else
	{
		_deadtimer.cancel ();
		boost::asio::async_write ( _socket, 
								  _request,
								  boost::bind (
											   &WebpageHunter::handle_write_request,
											   this,
											   boost::asio::placeholders::error
											   ) );
		_deadtimer.expires_from_now ( boost::posix_time::seconds ( _deadseconds ) );
		_deadtimer.async_wait ( boost::bind ( &WebpageHunter::handle_timeout, this, boost::asio::placeholders::error ) );

	}
}

void WebpageHunter::handle_write_request( const boost::system::error_code& err )
{
	if ( err )
	{
		_deadtimer.cancel();
		throw boost::system::system_error ( err ); 	
	}
	else
	{
		_deadtimer.cancel();
		boost::asio::async_read_until ( _socket,
									   _response,
									   "\r\n",
									   boost::bind ( 
													&WebpageHunter::handle_read_status_line,
													this,
													boost::asio::placeholders::error
													) );
		_deadtimer.expires_from_now ( boost::posix_time::seconds ( _deadseconds ) );
		_deadtimer.async_wait ( boost::bind ( &WebpageHunter::handle_timeout, this, boost::asio::placeholders::error ) );
	}
}

void WebpageHunter::handle_read_status_line( const boost::system::error_code& err )
{
	if ( err )
	{
		_deadtimer.cancel();
		throw boost::system::system_error ( err ); 	
	}
	else
	{
		_response_stream = new istream ( &_response );
		
		string http_version;
		*_response_stream >> http_version;
		
		unsigned int status_code;
		*_response_stream >> status_code;
		
		string status_message;
		getline ( *_response_stream, status_message );
		if ( !(*_response_stream) || http_version.substr ( 0,5 ) != "HTTP/" )
		{
			throw runtime_error ( "Invalid Response Header." );
		}
		
		if ( status_code != 200 )
		{
			throw runtime_error ( "Not Expected Status Code." );
		}
		
		_deadtimer.cancel();
		boost::asio::async_read_until ( _socket,
									   _response, 
									   "\r\n\r\n",
									   boost::bind ( 
													&WebpageHunter::handle_read_headers,
													this,
													boost::asio::placeholders::error
													) );
		_deadtimer.expires_from_now ( boost::posix_time::seconds ( _deadseconds ) );
		_deadtimer.async_wait ( boost::bind ( &WebpageHunter::handle_timeout, this, boost::asio::placeholders::error ) );
	}
}

void WebpageHunter::handle_read_headers ( const boost::system::error_code& err )
{
	if ( err )
	{
		_deadtimer.cancel();
		throw boost::system::system_error ( err );
	}
	else
	{
		
		string header;
		while ( getline ( *_response_stream, header ) && header !="\r" )
		{
			*_header_output << header << "\n";
		}
		
		*_header_output << "\n";
		
		if ( _response.size() > 0 )
		{
			*_content_output << &_response;
		}
		_deadtimer.cancel();
		boost::asio::async_read ( _socket,
								 _response, 
								 boost::asio::transfer_at_least(1),
								 boost::bind ( 
											  &WebpageHunter::handle_read_content,
											  this,
											  boost::asio::placeholders::error
											  ) );
		_deadtimer.expires_from_now ( boost::posix_time::seconds ( _deadseconds ) );
		_deadtimer.async_wait ( boost::bind ( &WebpageHunter::handle_timeout, this, boost::asio::placeholders::error ) );
	}
}

void WebpageHunter::handle_read_content ( const boost::system::error_code& err )
{
	if ( err )
	{
		if ( err == boost::asio::error::eof )
		{
			if ( _response.size() > 0 )
			{
				
				*_content_output << &_response;
			}
			
			_deadtimer.cancel();
		}
		else
		{
			_deadtimer.cancel();
			throw boost::system::system_error ( err );
		}
	}
	else
	{
		*_content_output << &_response;
		_deadtimer.cancel();
		boost::asio::async_read ( _socket,
								 _response, 
								 boost::asio::transfer_at_least(1),
								 boost::bind ( 
											  &WebpageHunter::handle_read_content,
											  this,
											  boost::asio::placeholders::error
											  ) );
		_deadtimer.expires_from_now ( boost::posix_time::seconds ( _deadseconds ) );
		_deadtimer.async_wait ( boost::bind ( &WebpageHunter::handle_timeout, this, boost::asio::placeholders::error ) );
	}
}

void WebpageHunter::handle_timeout ( const boost::system::error_code& err )
{
	if ( err == boost::asio::error::operation_aborted )
	{
		// do something after operation cancelled.
	}
	else
	{
		_socket.close();
	}
}

#if (defined __WEBPAGE_HUNTER_UNITTEST__)

int main ( int argc, char* argv[] )
{
	boost::asio::io_service io_service;
	
	HttpHeaderPairs headers;
	
	WebpageHunter h ( io_service, &cout, &cout, "www.sina.com.cn", "http", "www.sina.com.cn", "/", headers );
	io_service.run();
	
	return 0;
}

#endif